import platgo as pg
import numpy as np


class LVIDE(pg.Algorithm):

    def __init__(self, maxgen: int, problem: pg.Problem) -> None:
        super().__init__(maxgen=maxgen, problem=problem)
        self.name = "LVIDE"
        self.xov = pg.operators.XovSbx()  # 模拟二进制交叉
        self.mut = pg.operators.MutPol(problem)  # 多项式变异
        self.FE_max = 10 ** 6

    def go(self, N: int = None, population: pg.Population = None):
        assert N or population, "N and population can't be both None"
        if population is None:
            pop = self.problem.init_pop(N)
        else:
            pop = population
            self.problem.N = pop.decs.shape[0]
        self.problem.cal_obj(pop)
        F = 0.5
        if self.problem.D <= 1000:
            M = self.problem.D / 25
            u = 0.05 * (1 + self.problem.D / 100)
            K = 25
        else:
            M = self.problem.D / 50
            u = 0.6
            K = 15
        FE = 0
        # weight_vector, N = pg.utils.uniform_point(N=pop.N, M=self.problem.M, method='MUD')  # 生成权重向量
        # self.IVG(K, weight_vector, pop)
        W = [1 / self.problem.M] * self.problem.M
        while self.not_terminal(pop):
            G1, G2, dk = self.IVG(pop)
            print(G1)
            print(G2)
            break
        return pop

    def IVG(self, pop):
        """lb = self.problem.borders[0][0]
        ub = self.problem.borders[1][0]
        v = np.linspace(lb, ub, K)"""
        # temp = pg.Population(decs=pop.decs)
        # temp_f = []
        # for i in range(self.problem.D):
        #     temp_f1 = []
        #     for j in range(K):
        #         temp.decs[:, i:i+1] = v[j]
        #         self.problem.cal_obj(temp)
        #         temp_f1.append(temp.objv)
        #     print(np.array(temp_f1))
        #     break
        dk = self.FiniteDifference(pop)  # 求出种群所有个体在每个决策变量上梯度的和
        frontno, maxfront = pg.utils.nd_sort(np.minimum(dk, -dk), self.problem.D)
        temp = []
        G1 = []
        G2 = []
        #  分组
        for i in range(int(maxfront)):
            g1 = np.array(np.where(frontno == i + 1))
            temp += list(g1[0])
            G1.append(list(g1[0]))
            G2.append(list(temp))
        return G1, G2, dk

    def VIDE(self, G1, G2, pop, FE, FE_max):
        return

    def Environmental_Selection(self, pop, offspring):
        return

    def FiniteDifference(self, pop: pg.Population) -> np.ndarray:
        # Estimate the gradient of objective by finite difference
        gradient = []
        '''
        for i in range(pop.decs.shape[0]):
            temp = pg.Population(decs=pop.decs[i] + np.eye(pop.decs.shape[1]) * 1e-4)
            self.problem.cal_obj(temp)
            gradient.append(np.sum(((temp.objv - pop.objv[i]) / 1e-4) * W, axis=1))
        gradient = np.array(gradient)
        '''
        for i in range(pop.decs.shape[1]):
            temp = pop.copy()
            temp.decs[:, i:i + 1] += 1e-6
            self.problem.cal_obj(temp)
            gradient.append(np.sum((temp.objv - pop.objv) / 1e-6, axis=0))
        return np.array(gradient)


if __name__ == '__main__':
    problem = pg.problems.DTLZ1()
    N = 100
    maxgen = 50
    Algorithm = pg.algorithms.LVIDE(maxgen=maxgen, problem=problem)
    pop = Algorithm.go(N)
